package com.tomas.demo.common;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

import com.alibaba.fastjson.JSONObject;
import com.tomas.demo.common.cache.Cacher;
import com.tomas.demo.common.utils.IpUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.HandlerInterceptor;
import org.springframework.web.servlet.ModelAndView;

import java.io.IOException;
import java.io.PrintWriter;

/**
 * @author  Tomas
 * @date 2016-04-11 18:00:00
 * @what URL防刷拦截器 限制单个IP每秒种请求数量 requests/seconds
 */
@Component
public class RateLimitInterceptor implements HandlerInterceptor {
    private static final Logger logger = LoggerFactory.getLogger(RateLimitInterceptor.class);
    @Autowired
    private Cacher cacher;
    @Value("${spring.redis.api.ratelimit:300}")
    private Long ratelimit;

    private final String ipPrefix = "::ip::";

    public RateLimitInterceptor() {
    }

    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        if(this.ratelimit != null) {
            String ip = IpUtils.getIpAddress(request);
            if(ip != null && !ip.isEmpty() && !"unknown".equalsIgnoreCase(ip)) {
                Long expire = this.cacher.getExpire("::ip::" + ip);
                if(expire.longValue() != 1L) {
                    this.cacher.expire("::ip::" + ip, 0L);
                    this.cacher.remove("::ip::" + ip);
                }

                Long rate = this.cacher.increment("::ip::" + ip, 1L, -1L);
                if(rate.longValue() == 1L) {
                    this.cacher.expire("::ip::" + ip, 1L);
                } else if(rate.longValue() >= this.ratelimit.longValue()) {
                    response.setContentType("application/json;charset=UTF-8");
                    try {
                        PrintWriter writer = response.getWriter();
                        JSONObject content=new JSONObject();
                        content.put("code",40);
                        content.put("message","请求次数超过限制");
                        writer.print(content);
                        writer.flush();
                    } catch (IOException e) {
                        logger.error(e.getMessage(), e);
                    }
                    logger.warn("Too many requests per seconds for ip \'" + ip + "\', rate limit " + this.ratelimit + ", requests " + (rate.longValue() + 1L) + ".");
                    return false;
                }
            }
        }

        return true;
    }

    public void postHandle(HttpServletRequest request, HttpServletResponse response, Object handler, ModelAndView modelAndView) throws Exception {
    }

    public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler, Exception ex) throws Exception {
    }
}
